# 定义LLM
import os

from langchain.agents import ZeroShotAgent, AgentExecutor
from langchain.chat_models import init_chat_model
from langchain_core.prompts import PromptTemplate
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI

# 定义LLM--》大脑
llm = ChatOpenAI(
    model_name="qwen-plus",
    api_key=os.getenv("DASHSCOPE_API_KEY"),
    temperature=0,
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1"
)


# 定义工具
@tool(
    description="用于进行加减乘除运算,参数:expression(数学表达式,必填,示例:1+1,2*4),返回值:运算结果"
)
def calculatorTool(expression: str) -> str:
    res = eval(expression)
    return res


# 关键步骤:把工具放到集合里面
tools = [calculatorTool]

prompt_template = """
你是一个帮助用户解决问题的智能助手，你拥有以下工具:
{tools}

使用规则:
1:先理解问题:无需工具则直接回答,需要工具则调用工具(根据工具的描述匹配)
2:调用工具必须用格式：```json{{"name":"工具名","parameters":{{"参数名":"参数值"}}}}```（参数需齐全）；
3:拿到工具结果后, 判断是否继续调用工具, 如果需要继续调用工具,则调用工具,否则整理,返回结果;


用户的问题:{input}
"""

prompt = PromptTemplate(
    template=prompt_template,
    input_variables=["tools", "input"],
    partial_variables={
        "tools": "\n".join([f"{t.name}:{t.description}" for t in tools])
    }
)

# 创建agent: 零样本类型
agent = ZeroShotAgent.from_llm_and_tools(
    llm=llm,
    tools=tools,
    prompt=prompt,
    verbose=True
)

executor = AgentExecutor(
    agent=agent,
    tools=tools,
    verbose=True,
    handle_parsing_errors=True,
    max_iterations=5,
)

if __name__ == '__main__':
    res = executor.invoke({
        "input": "1+89等于什么?"
    })
    print(res["output"])
